import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import cv2


class RectScale(object):
    def __init__(self, height, width, interpolation=Image.BILINEAR):
        self.height = height
        self.width = width
        self.interpolation = interpolation

    def __call__(self, img):
        w, h = img.size
        if h == self.height and w == self.width:
            return img
        return img.resize((self.width, self.height), self.interpolation)


personTransform1 = transforms.Compose([
    RectScale(384, 128),
])

personTransform2 = transforms.Compose([
    transforms.Resize((256, 128)),
])


class PersonTransform3:
    def __init__(self):
        self.transform3 = transforms.Compose([
            transforms.Resize((256, 128)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
            # transforms.ToPILImage(),
        ])

    def __call__(self, x):
        t_img = self.transform3(x)
        img_ = t_img.numpy()
        img_ = np.transpose(img_, (1, 2, 0))
        img_ = img_*255.
        img_ = img_[:, :, ::-1]
        cv2.imwrite('temp_pic.jpg', img_)
        img = Image.open('temp_pic.jpg')
        return img


personTransform3 = PersonTransform3()


class PersonTransform4:
    def __init__(self):
        self.transform4 = transforms.Compose([
            # transforms.Resize((256, 128)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
            # transforms.ToPILImage(),
        ])

    def __call__(self, x):
        t_img = self.transform4(x)
        img_ = t_img.numpy()
        img_ = np.transpose(img_, (1, 2, 0))
        img_ = img_*255.
        img_ = img_[:, :, ::-1]
        cv2.imwrite('temp_pic.jpg', img_)
        img = Image.open('temp_pic.jpg')
        return img


personTransform4 = PersonTransform4()


if __name__ == '__main__':
    import cv2
    img = Image.open('img1.jpg').convert("RGB")
    img = personTransform3(img)
    img.save("img1_new.jpg")

    # img = Image.open('img.jpg').convert("RGB")
    # img = personTransform4(img)
    # img.save("img_new.jpg")
